# Copyright (c) 2010
# The Regents of the University of Michigan
# All Rights Reserved

# Permission is granted to use, copy, create derivative works, and
# redistribute this software and such derivative works for any purpose,
# so long as the name of the University of Michigan is not used in
# any advertising or publicity pertaining to the use or distribution
# of this software without specific, written prior authorization. If
# the above copyright notice or any other identification of the
# University of Michigan is included in any copy of any portion of
# this software, then the disclaimer below must also be included.

# This software is provided as is, without representation or warranty
# of any kind either express or implied, including without limitation
# the implied warranties of merchantability, fitness for a particular
# purpose, or noninfringement.  The Regents of the University of
# Michigan shall not be liable for any damages, including special,
# indirect, incidental, or consequential damages, with respect to any
# claim arising out of or in connection with the use of the software,
# even if it has been or is hereafter advised of the possibility of
# such damages.


# This file is derived from ktable.py in the Khashmir project
# (http://khashmir.sourceforge.net/), whose copyright notice is as
# follows:

# Copyright (C) 2002-2003 Andrew Loewenstern

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:

# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

# The Software is provided "AS IS", without warranty of any kind,
# express or implied, including but not limited to the warranties of
# merchantability,  fitness for a particular purpose and
# noninfringement. In no event shall the  authors or copyright holders
# be liable for any claim, damages or other liability, whether in an
# action of contract, tort or otherwise, arising from, out of or in
# connection with the Software or the use or other dealings in the
# Software.

import bisect
import cPickle as pickle
import functools
import time

from pydht import contact, dht, dhtutil
from dhtutil import intify


MAX_FAILURES = 2
K = 20
HASH_LENGTH = 160  # SHA-1 is 20 bytes, right?


class RoutingTableError(Exception):
  pass


class DummyRoutingTable(object):
  def get_k_closest(self, node_id):
    return []

  def save(self, fyl):
    pass

  def load(self, fyl):
    pass

  def replace_stale_node(self, stale, new):
    pass

  def record_live(self, node, live=True):
    pass

  def remove_node(self, node):
    pass

  def didnt_respond(self, node):
    pass


# TODO TODO TODO: should this module hide node state required to
# decide whether a node should remain in the routing table, or expose
# it? The former makes life simpler when you don't want to use the
# routing table for some reason (e.g., you have global routing
# information), but the latter makes it possible to use said
# information. Maybe it's possible to add a wrapper or a
# "get_routing_info(contact)" function?


class RoutingTable(object):
  '''Kademlia routing table.

  Port of ktable from Khashmir.
  '''
  def __init__(self, local_contact):
    self.local_contact = contact.Contact(local_contact)
    self.buckets = [Bucket([], 0, 2**HASH_LENGTH)]
    self.record_live(local_contact)

  def save(self, fyl):
    pickle.dump((self.local_contact, self.buckets), fyl, pickle.HIGHEST_PROTOCOL)

  def load(self, fyl):
    (lc, buckets) = pickle.load(fyl)
    if lc.addrport != self.local_contact.addrport:
      err = 'saved routing table IP/port %s:%d does not match %s:%d'
      err = err % (lc.addrport + self.local_contact.addrport)
      raise RoutingTableError(err)
    self.buckets = buckets
    # TODO: scan buckets for contact liveness, update appropriately.

  def _bucket_index_for_node(self, n):
    return bisect.bisect_left(self.buckets, n)

  def get_k_closest(self, node_id, exact=False):
    if isinstance(node_id, str):
      node_id = intify(node_id)
    elif isinstance(node_id, tuple):
      node_id = intify(dht.compute_node_id(node_id))
    elif isinstance(node_id, contact.Contact):
      node_id = intify(node_id.dht_id)

    assert isinstance(node_id, long)
    bucket_idx = self._bucket_index_for_node(node_id)

    bucket = self.buckets[bucket_idx]
    try:
      idx = bucket.contents.index(node_id)
      if exact:  # exact is for internal use only.
        return [bucket.contents[idx]]
      else:
        return [bucket.contents[idx].addrport]
    except ValueError:
      pass

    if exact:
      return None

    nodes = bucket.contents[:]  # Don't go handing out/modifying the bucket!
    if len(nodes) < K:
      min = bucket_idx - 1
      max = bucket_idx + 1
      nbuckets = len(self.buckets)
      # Crawl outward from the target bucket.
      while len(nodes) < K and (min >= 0 or max < nbuckets):
        # Note ported from Khashmir: this requires that K is even.
        if min >= 0:
          nodes.extend(self.buckets[min].contents)
        if max < nbuckets:
          nodes.extend(self.buckets[max].contents)
        min -= 1
        max += 1

    if len(nodes) != K:
      nodes.sort(cmp=lambda a, b: cmp(node_id ^ intify(a.dht_id),
                                      node_id ^ intify(b.dht_id)))

    return [x.addrport for x in nodes[:K]]

  def _split_bucket(self, bucket):
    diff = (bucket.max - bucket.min) / 2
    new = Bucket([], bucket.max - diff, bucket.max)
    self.buckets.insert(self.buckets.index(bucket.min) + 1, new)
    bucket.max -= diff

    # Original code does remove and append, which makes it quadratic
    # in the number of removed items. We do more passes, but linear
    # time.
    to_remove = [False] * len(bucket.contents)
    for idx, node in enumerate(bucket.contents):
      if intify(node.dht_id) > bucket.max:
        to_remove[idx] = True

    zipped = zip(bucket.contents, to_remove)
    bucket.contents = [x for x, removed in zipped if not removed]
    new.contents = [x for x, removed in zipped if removed]

  def replace_stale_node(self, stale, new):
    stale_id = intify(stale.dht_id)
    bucket_idx = self._bucket_index_for_node(stale_id)
    try:
      victim_idx = self.buckets[bucket_idx].contents.remove(stale_id)
    except ValueError:
      return

    if new:
      self.buckets[bucket_idx].contents.append(new)

  def record_live(self, node, live=True):
    '''Insert node into the routing table, unless already there.
    Mark it live in either case.
    '''
    if not isinstance(node, contact.Contact):
      node = contact.Contact(node)
    if node.dht_id == self.local_contact.dht_id:
      return

    node_id = intify(node.dht_id)
    bucket = self.buckets[self._bucket_index_for_node(node_id)]
    try:
      # Is it in the bucket?
      node_idx = bucket.contents.index(node_id)
    except ValueError:
      pass
    else:
      node = bucket.contents[node_idx]
      if live:
        node.seen()
        # Move to end of bucket.
        del bucket.contents[node_idx]
        bucket.contents.append(node)
        bucket.touch()
      return


    # If bucket isn't full, stick the new node in the bucket.
    if len(bucket.contents) < K:
      if live:
        node.seen()
      bucket.contents.append(node)
      bucket.touch()
      return

    # Bucket's full.
    
    # If our node isn't in the bucket, return the oldest node in it,
    # so higher layers can think about replacing said node. We can't
    # do it here because that would completely force the caller to
    # confirm to whatever event model we chose. If you want it done
    # here, it shouldn't be too hard to patch in.

    # TODO:
    # Make said higher layers sane, think about what this class
    # encapsulates, etc.
    #
    # Implement splitting buckets that don't contain our node, as
    # described in Kademlia paper.
    if self.local_contact != bucket:  # Bucket overrides __ne__, remember.
      oldie = bucket.contents[0]
      if live:
        bucket.contents = bucket.contents[1:]
        bucket.contents.append(node)
        node.seen()
      else:
        return bucket.contents[0]

    # Can't split the bucket, table's full.
    if len(self.buckets) >= HASH_LENGTH:
      return

    self._split_bucket(bucket)

    # Tail recursion is OK, bucket splits don't happen often.
    return self.record_live(node)

  def remove_node(self, node):
    '''Forcibly remove a node from the routing table.'''
    self.replace_stale_node(node, None)

  def didnt_respond(self, node):
    '''Record that node did not respond to a message.'''
    try:
      n = self.get_k_closest(node, exact=True)[0]
    except (IndexError, TypeError):
      return
    if n.didnt_respond() >= MAX_FAILURES:
      self.remove_node(n)


# @functools.total_ordering  # oops, python 3 only.
class Bucket(object):
  __slots__ = ('min', 'max', 'lastAccessed', 'contents')
  def __init__(self, contents, min, max):
    self.contents = contents
    self.min = min
    self.max = max
    self.lastAccessed = time.time()

  def touch(self):
    self.lastAccessed = time.time()

  def __repr__(self):
    return 'Bucket(%r, %r, %r)' % (self.contents, self.min, self.max) 

  def __lt__(self, rhs):
    if isinstance(rhs, tuple):
      rhs = intify(dht.compute_node_id(rhs))
    elif isinstance(rhs, contact.Contact):
      rhs = intify(rhs.dht_id)

    assert isinstance(rhs, long)
    return self.max <= rhs

  def __le__(self, rhs):
    if isinstance(rhs, tuple):
      rhs = intify(dht.compute_node_id(rhs))
    elif isinstance(rhs, contact.Contact):
      rhs = intify(rhs.dht_id)
    assert isinstance(rhs, long)
    return self.min <= rhs

  def __gt__(self, rhs):
    if isinstance(rhs, tuple):
      rhs = intify(dht.compute_node_id(rhs))
    elif isinstance(rhs, contact.Contact):
      rhs = intify(rhs.dht_id)
    assert isinstance(rhs, long)
    return self.min > rhs

  def __ge__(self, rhs):
    if isinstance(rhs, tuple):
      rhs = intify(dht.compute_node_id(rhs))
    elif isinstance(rhs, contact.Contact):
      rhs = intify(rhs.dht_id)
    assert isinstance(rhs, long)
    return self.max > rhs 
  
  def __eq__(self, rhs):
    if isinstance(rhs, tuple):
      rhs = intify(dht.compute_node_id(rhs))
    elif isinstance(rhs, contact.Contact):
      rhs = intify(rhs.dht_id)

    assert isinstance(rhs, long) or isinstance(rhs, int)
    return self.min <= rhs < self.max

  def __ne__(self, rhs):
    return not self == rhs


class RouterAwareClient(dht.DHTClient):
  # XXX: assumes client implementation uses _process_packet or similar.
  def set_router(self, router):
    self.router = router

  def _process_packet(self, packet, sender):
    if dht.DHTClient._process_packet(self, packet, sender):
      self.router.record_live(contact.Contact(sender))
      return True
    return False
